{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "15 - Distributed embeddings cluster",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4Pjmz-RORV8E"
      },
      "source": [
        "# Distributed embeddings cluster\n",
        "\n",
        "_This notebook is part of a tutorial series on [txtai](https://github.com/neuml/txtai), an AI-powered semantic search platform._\n",
        "\n",
        "The txtai API is a web-based service backed by [FastAPI](https://fastapi.tiangolo.com/). All txtai functionality is available via the API. The API can also cluster multiple embeddings indices into a single logical index to horizontally scale over multiple nodes. \n",
        "\n",
        "This notebook installs the txtai API and shows an example of building an embeddings cluster."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Dk31rbYjSTYm"
      },
      "source": [
        "# Install dependencies\n",
        "\n",
        "Install `txtai` and all dependencies. Since this notebook uses the API, we need to install the api extras package."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XMQuuun2R06J"
      },
      "source": [
        "%%capture\n",
        "!pip install git+https://github.com/neuml/txtai#egg=txtai[api]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PNPJ95cdTKSS"
      },
      "source": [
        "# Start distributed embeddings cluster\n",
        "\n",
        "First we'll start multiple API instances that will serve as embeddings index shards. Each shard stores a subset of the indexed data and these shards work in tandem to form a single logical index.\n",
        "\n",
        "Then we'll start the main API instance that clusters the shards together into a logical instance.\n",
        "\n",
        "The API instances are all started in the background.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "USb4JXZHxqTA"
      },
      "source": [
        "import os\n",
        "os.chdir(\"/content\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nTDwXOUeTH2-",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "dee26849-39ae-4390-8bba-76bf9025fa61"
      },
      "source": [
        "%%writefile index.yml\n",
        "writable: true\n",
        "\n",
        "# Embeddings settings\n",
        "embeddings:\n",
        "    path: sentence-transformers/nli-mpnet-base-v2\n",
        "    content: true"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Writing index.yml\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "iCdBh-JgfyBl",
        "outputId": "0066e314-7461-47c7-ca3b-15204911783e"
      },
      "source": [
        "%%writefile cluster.yml\n",
        "# Embeddings cluster\n",
        "cluster:\n",
        "    shards:\n",
        "        - http://127.0.0.1:8001\n",
        "        - http://127.0.0.1:8002"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Writing cluster.yml\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nGITHxUyRzyp"
      },
      "source": [
        "# Start embeddings shards\n",
        "!CONFIG=index.yml nohup uvicorn --port 8001 \"txtai.api:app\" &> shard-1.log &\n",
        "!CONFIG=index.yml nohup uvicorn --port 8002 \"txtai.api:app\" &> shard-2.log &\n",
        "\n",
        "# Start main instance\n",
        "!CONFIG=cluster.yml nohup uvicorn --port 8000 \"txtai.api:app\" &> main.log &\n",
        "\n",
        "# Wait for startup\n",
        "!sleep 90"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lxkbVng3giWP"
      },
      "source": [
        "# Python\n",
        "\n",
        "Let's first try the cluster out directly in Python. The code below aggregates the two shards into a single cluster and executes actions against the cluster."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "36HGAokoglfg",
        "outputId": "368ae013-2afc-4a1b-d7df-c429183637d7"
      },
      "source": [
        "%%writefile run.py\n",
        "from txtai.api import Cluster\n",
        "\n",
        "cluster = Cluster({\"shards\": [\"http://127.0.0.1:8001\", \"http://127.0.0.1:8002\"]})\n",
        "\n",
        "data = [\n",
        "    \"US tops 5 million confirmed virus cases\",\n",
        "    \"Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg\",\n",
        "    \"Beijing mobilises invasion craft along coast as Taiwan tensions escalate\",\n",
        "    \"The National Park Service warns against sacrificing slower friends in a bear attack\",\n",
        "    \"Maine man wins $1M from $25 lottery ticket\",\n",
        "    \"Make huge profits without work, earn up to $100,000 a day\",\n",
        "]\n",
        "\n",
        "# Index data\n",
        "cluster.add([{\"id\": x, \"text\": row} for x, row in enumerate(data)])\n",
        "cluster.index()\n",
        "\n",
        "# Test search\n",
        "result = cluster.search(\"feel good story\", 1)[0]\n",
        "print(\"Query: feel good story\\nResult:\", result[\"text\"])"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Writing run.py\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6dQOzcfEs2Pk",
        "outputId": "a667594a-b778-4e4e-a75c-72e7982b7fbe"
      },
      "source": [
        "!python run.py"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Query: feel good story\n",
            "Result: Maine man wins $1M from $25 lottery ticket\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NHvBFZeSd9AG"
      },
      "source": [
        "# JavaScript\n",
        "\n",
        "Next let's try to run the same code above via the API using JavaScript.\n",
        "\n",
        "```bash\n",
        "npm install txtai\n",
        "```\n",
        "\n",
        "For this example, we'll clone the txtai.js project to import the example build configuration."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "b52knObEdcCr"
      },
      "source": [
        "%%capture\n",
        "!git clone https://github.com/neuml/txtai.js"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rUGS0t-JMsS9"
      },
      "source": [
        "## Run cluster.js\n",
        "\n",
        "The following script is a JavaScript version of the logic above"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "bPQ40_xRyFmA",
        "outputId": "b86a12c4-f2c7-427b-bd28-edba354c6713"
      },
      "source": [
        "%%writefile txtai.js/examples/node/src/cluster.js\n",
        "import {Embeddings} from \"txtai\";\n",
        "import {sprintf} from \"sprintf-js\";\n",
        "\n",
        "const run = async () => {\n",
        "    try {\n",
        "        let embeddings = new Embeddings(process.argv[2]);\n",
        "\n",
        "        let data  = [\"US tops 5 million confirmed virus cases\",\n",
        "                     \"Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg\",\n",
        "                     \"Beijing mobilises invasion craft along coast as Taiwan tensions escalate\",\n",
        "                     \"The National Park Service warns against sacrificing slower friends in a bear attack\",\n",
        "                     \"Maine man wins $1M from $25 lottery ticket\",\n",
        "                     \"Make huge profits without work, earn up to $100,000 a day\"];\n",
        "\n",
        "        console.log();\n",
        "        console.log(\"Querying an Embeddings cluster\");\n",
        "        console.log(sprintf(\"%-20s %s\", \"Query\", \"Best Match\"));\n",
        "        console.log(\"-\".repeat(50));\n",
        "\n",
        "        for (let query of [\"feel good story\", \"climate change\", \"public health story\", \"war\", \"wildlife\", \"asia\", \"lucky\", \"dishonest junk\"]) {\n",
        "            let results = await embeddings.search(query, 1);\n",
        "            if (results && results.length > 0) {\n",
        "              let result = results[0].text;\n",
        "              console.log(sprintf(\"%-20s %s\", query, result));\n",
        "            }\n",
        "        }\n",
        "    }\n",
        "    catch (e) {\n",
        "        console.trace(e);\n",
        "    }\n",
        "};\n",
        "\n",
        "run();"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Writing txtai.js/examples/node/src/cluster.js\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nTBs11j-GtD-"
      },
      "source": [
        "## Build and run cluster.js\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kC5Oub6wa1nK"
      },
      "source": [
        "%%capture\n",
        "os.chdir(\"txtai.js/examples/node\")\n",
        "!npm install\n",
        "!npm run build"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Xr5IlvqH8W77"
      },
      "source": [
        "Next lets run the code against the main cluster URL"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ckOHNqyaeL-B",
        "outputId": "9c243fac-2316-4b8e-b044-6de529a8f3e8"
      },
      "source": [
        "!node dist/cluster.js http://127.0.0.1:8000"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Querying an Embeddings cluster\n",
            "Query                Best Match\n",
            "--------------------------------------------------\n",
            "feel good story      Maine man wins $1M from $25 lottery ticket\n",
            "climate change       Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg\n",
            "public health story  US tops 5 million confirmed virus cases\n",
            "war                  Beijing mobilises invasion craft along coast as Taiwan tensions escalate\n",
            "wildlife             The National Park Service warns against sacrificing slower friends in a bear attack\n",
            "asia                 Beijing mobilises invasion craft along coast as Taiwan tensions escalate\n",
            "lucky                Maine man wins $1M from $25 lottery ticket\n",
            "dishonest junk       Make huge profits without work, earn up to $100,000 a day\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1yukBIMYG5OE"
      },
      "source": [
        "The JavaScript program is showing the same results as the Python code above. This is running a clustered query against both nodes in the cluster and aggregating the results together.\n",
        "\n",
        "Queries can be run against each individual shard to see what the queries independently return."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "73rZCo4O4IQR",
        "outputId": "9f2cb119-7a21-41d9-fdbf-4410af246934"
      },
      "source": [
        "!node dist/cluster.js http://127.0.0.1:8001"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Querying an Embeddings cluster\n",
            "Query                Best Match\n",
            "--------------------------------------------------\n",
            "feel good story      Maine man wins $1M from $25 lottery ticket\n",
            "climate change       Beijing mobilises invasion craft along coast as Taiwan tensions escalate\n",
            "public health story  US tops 5 million confirmed virus cases\n",
            "war                  Beijing mobilises invasion craft along coast as Taiwan tensions escalate\n",
            "wildlife             Beijing mobilises invasion craft along coast as Taiwan tensions escalate\n",
            "asia                 Beijing mobilises invasion craft along coast as Taiwan tensions escalate\n",
            "lucky                Maine man wins $1M from $25 lottery ticket\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ZeVBLJyr4Knr",
        "outputId": "b75691a4-25bf-43dc-8878-f9792a4430b8"
      },
      "source": [
        "!node dist/cluster.js http://127.0.0.1:8002"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Querying an Embeddings cluster\n",
            "Query                Best Match\n",
            "--------------------------------------------------\n",
            "feel good story      Make huge profits without work, earn up to $100,000 a day\n",
            "climate change       Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg\n",
            "public health story  The National Park Service warns against sacrificing slower friends in a bear attack\n",
            "war                  The National Park Service warns against sacrificing slower friends in a bear attack\n",
            "wildlife             The National Park Service warns against sacrificing slower friends in a bear attack\n",
            "asia                 The National Park Service warns against sacrificing slower friends in a bear attack\n",
            "lucky                The National Park Service warns against sacrificing slower friends in a bear attack\n",
            "dishonest junk       Make huge profits without work, earn up to $100,000 a day\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J2I_4hmZ8uXs"
      },
      "source": [
        "Note the differences. The section below runs a count against the full cluster and each shard to show the count of records in each."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "BKm27yna4MWr",
        "outputId": "bfc60af7-1b2b-451f-b10e-e2f8cf6f14fa"
      },
      "source": [
        "!curl http://127.0.0.1:8000/count\n",
        "!printf \"\\n\"\n",
        "!curl http://127.0.0.1:8001/count\n",
        "!printf \"\\n\"\n",
        "!curl http://127.0.0.1:8002/count"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "6\n",
            "3\n",
            "3"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6rKj-I0djRQj"
      },
      "source": [
        "This notebook showed how a distributed embeddings cluster can be created with txtai. This example can be further scaled out on Kubernetes with StatefulSets, which will be covered in a future tutorial."
      ]
    }
  ]
}